# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time : 2020/4/9 14:00
@Author : SPZ
@File : data_loader.py
@Software: pycharm
"""
from concurrent.futures import ThreadPoolExecutor
try:
    import moxing as mox
except Exception as info:
    str(info)
import numpy as np
import os
import cv2


class _RandomUniformSelector:
    def __init__(self, total, start=0):
        self.total = total
        self.idxes = [i for i in range(total)]
        np.random.shuffle(self.idxes)
        self.cursor = start

    def next(self, batch_size):
        if self.cursor + batch_size > self.total:
            r_n = []
            r_n_1 = self.idxes[self.cursor:self.total]
            self.cursor = self.cursor + batch_size - self.total
            r_n_2 = self.idxes[0:self.cursor]
            r_n.extend(r_n_1)
            r_n.extend(r_n_2)
        else:
            r_n = self.idxes[self.cursor:self.cursor + batch_size]
            self.cursor = self.cursor + batch_size
        return r_n


class DataLoader:

    def __init__(self, dict_file_path,
                 labeled_file_path,
                 images_dir,
                 encoding='utf-8',
                 blank_first=False,
                 batch_size=64,
                 maxlen=10,
                 image_shape=(32, 280),
                 **kwgs
                 ):
        self.images_dir = images_dir
        self.blank_first = blank_first
        self.maxlen = maxlen
        self.image_shape = image_shape
        self.id_to_dict = self.__load_dict(dict_file_path, encoding)
        self.num_classes = len(self.id_to_dict)
        self.batch_size = batch_size
        self.image_label = self.__load_labeled_file(labeled_file_path, encoding)
        self.image_files = list(self.image_label.keys())
        self.total_size = len(self.image_files)
        self.random_uniform_selector = _RandomUniformSelector(self.total_size)
        self.steps_per_epoch = kwgs['steps_per_epoch'] if 'steps_per_epoch' in kwgs else self.total_size // self.batch_size

    def __load_dict(self, dict_file_path, encoding='utf-8'):
        with open(dict_file_path, encoding=encoding, mode='r') as f:
            chars = list(map(lambda char: char.strip('\r\n'), f.readlines()))

        if self.blank_first:
            chars = chars[1:] + ['blank']

        return {i: v for i, v in enumerate(chars)}

    def __load_labeled_file(self, labeled_file_path, encoding='utf-8'):
        dic = {}
        with open(labeled_file_path, encoding=encoding, mode="r") as f:
            for line in f.readlines():
                line = line.strip()
                line = line.split(' ')
                dic[line[0]] = line[1:11]
        return dic

    def load_data(self):
        image_files = np.array(self.image_files)

        def load_single_example(image_file_path, image_label):

            n, ext = os.path.splitext(image_file_path)
            try:
                with mox.file.File(os.path.join(self.images_dir, (n+'.jpg')), 'rb') as f:
                    file_str = f.read()
            except Exception as info:
                # print(info)
                with open(os.path.join(self.images_dir, (n+'.jpg')), 'rb') as f:
                    file_str = f.read()
            im = cv2.imdecode(np.fromstring(file_str, np.uint8), cv2.IMREAD_GRAYSCALE)
            img = im / 255.0 - 0.5
            img = np.expand_dims(img, axis=2)

            label_len = np.array([len(image_label)])
            input_len = np.array([self.image_shape[1] // 8])
            label = np.ones([self.maxlen]) * self.num_classes

            if self.blank_first:
                label[0: len(image_label)] = [int(i) - 1 for i in image_label]
            else:
                label[0: len(image_label)] = [int(i) for i in image_label]

            result_info = dict()
            result_info['img'] = img
            result_info['label_len'] = label_len
            result_info['input_len'] = input_len
            result_info['label'] = label
            return result_info

        while True:
            shuffled_image_files = image_files[self.random_uniform_selector.next(self.batch_size)]
            
            char_info = [(image_file, self.image_label[image_file]) for image_file in shuffled_image_files]
            imgs = []
            label_lens = []
            input_lens = []
            labels = []

            with ThreadPoolExecutor() as executor:
                for info in executor.map(lambda t: load_single_example(*t), char_info):
                    imgs.append(info['img'])
                    label_lens.append(info['label_len'])
                    input_lens.append(info['input_len'])
                    labels.append(info['label'][:10])
            
            inputs = {'the_input': np.array(imgs),
                      'the_labels': np.array(labels),
                      'input_length': np.array(input_lens),
                      'label_length': np.array(label_lens),
                      }
            outputs = {
                    'ctc': np.zeros([self.batch_size]),
                    }

            yield inputs, outputs


